import cv2
import numpy as np
import os
from PIL import ImageGrab

rows = 14
cols = 10


def segment_chessboard(image, rows, cols):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    _, binary = cv2.threshold(
        gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 1))
    binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)

    contours, _ = cv2.findContours(
        binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    contours = sorted(contours, key=lambda c: cv2.boundingRect(c)[1])

    chess_pieces = [[None for _ in range(cols)] for _ in range(rows)]

    for i, contour in enumerate(contours):
        x, y, w, h = cv2.boundingRect(contour)
        row = i // cols
        col = i % cols

        if 0 <= row < rows and 0 <= col < cols:
            chess_pieces[row][col] = image[y:y+h, x:x+w]
    return chess_pieces


output_dir = "fenge"
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

clipboard_image = np.array(ImageGrab.grabclipboard())

if clipboard_image is not None:
    chess_pieces_images = segment_chessboard(clipboard_image, rows, cols)

    for row_idx, row in enumerate(chess_pieces_images):
        for col_idx, img in enumerate(row):
            if img is not None:
                filename = f"chess_piece_{row_idx + 1}_{cols - col_idx}.jpg"
                cv2.imwrite(os.path.join(output_dir, filename), img)

    print(f"棋盘成功分割成 {rows * cols} 个棋子图像。")
else:
    print("无法从剪贴板读取图像。")

# 棋局匹配部分
if not os.path.exists('muban'):
    os.makedirs('muban')

templates = []
for i in range(1, 53):
    template_path = os.path.join('muban', f'template_{i}.jpg')
    if os.path.isfile(template_path):
        template = cv2.imread(template_path, cv2.IMREAD_GRAYSCALE)
        templates.append(template)

chess_board = [[0] * cols for _ in range(rows)]

for x in range(rows):
    for y in range(cols):
        image_path = os.path.join('fenge', f'chess_piece_{x+1}_{y+1}.jpg')
        if not os.path.isfile(image_path):
            print(f'图像文件不存在：{image_path}')
            continue
        image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
        if image is None:
            print(f'无法加载图像：{image_path}')
            continue

        # 初始化最佳匹配的模板和相似度
        best_match = 0
        best_similarity = 0

        # 遍历模板进行匹配
        for j, template in enumerate(templates):
            res = cv2.matchTemplate(image, template, cv2.TM_CCOEFF_NORMED)
            min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)

            if max_val > best_similarity and max_val >= 0.4:
                best_match = j + 1  # 棋子数字为模板的索引加1
                best_similarity = max_val

        chess_board[x][y] = best_match

with open('numbers.txt', 'w') as f:
    for row in chess_board:
        f.write(' '.join(str(num) for num in row) + '\n')

for row in chess_board:
    print(row)
